import torch
import torch.nn as nn
import math
import SpectFormer as SGN

# 定义自定义的 Transpose 模块，用于在 nn.Sequential 中进行张量维度交换
class Transpose(nn.Module):
    def __init__(self, dim0, dim1):
        super().__init__()
        self.dim0 = dim0  # 第一个要交换的维度
        self.dim1 = dim1  # 第二个要交换的维度

    def forward(self, x):
        # 在前向传播中交换指定的维度
        return x.transpose(self.dim0, self.dim1)



# 定义简单的视觉 Transformer（SimpleViT）
class SimpleViT(nn.Module):
    def __init__(
        self, image_size=224, patch_size=16, num_classes=1000,
        dim=768, depth=6, heads=12, mlp_dim=3072
    ):
        super().__init__()
        # 确保图像尺寸能被 patch 大小整除
        assert image_size % patch_size == 0, '图像尺寸必须能被 patch 大小整除。'
        num_patches = (image_size // patch_size) ** 2  # 计算图像块的数量

        # 计算频域中的高度和宽度维度
        self.h = image_size // patch_size          # h = 224 // 16 = 14
        self.w = self.h // 2 + 1                   # w = 14 // 2 + 1 = 8

        self.patch_size = patch_size  # 图像块的大小

        # 定义 Patch Embedding 模块，将图像划分为图像块并嵌入到特征向量中
        self.to_patch_embedding = nn.Sequential(
            nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size),
            nn.Flatten(2),  # 将空间维度展平成一个维度
            Transpose(1, 2),  # 交换维度，使形状变为 (B, N, dim)
        )

        # 定义频谱门控网络，增强特征表示
        self.spectral_gating = SGN.SpectralGatingNetwork(dim, h=self.h, w=self.w)

        # 定义位置嵌入和分类令牌（CLS Token）
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))

        # 定义 Transformer 编码器层
        self.transformer = nn.ModuleList([])
        for _ in range(depth):
            self.transformer.append(
                nn.ModuleDict(
                    {
                        'norm1': nn.LayerNorm(dim),
                        'attn': nn.MultiheadAttention(dim, heads),
                        'norm2': nn.LayerNorm(dim),
                        'mlp': nn.Sequential(
                            nn.Linear(dim, mlp_dim),
                            nn.GELU(),
                            nn.Linear(mlp_dim, dim),
                        ),
                    }
                )
            )

        # 定义分类头，将最终的特征映射到类别预测上
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes),
        )

    def forward(self, img):
        B, C, H, W = img.shape  # 获取输入图像的批次大小、通道数、高度和宽度
        x = self.to_patch_embedding(img)  # 将图像转换为图像块嵌入，形状为 (B, N, dim)
        N = x.shape[1]  # 获取图像块的数量

        # 应用频谱门控网络，增强特征表示
        x = self.spectral_gating(x)

        # 添加分类令牌（CLS Token）
        cls_tokens = self.cls_token.expand(B, -1, -1)  # 扩展到批次大小，形状为 (B, 1, dim)
        x = torch.cat((cls_tokens, x), dim=1)  # 将 CLS Token 添加到嵌入序列前面，形状为 (B, N+1, dim)

        # 添加位置嵌入，保留图像块的位置信息
        x += self.pos_embedding[:, : N + 1]

        # 通过 Transformer 编码器层，捕获全局特征关系
        for layer in self.transformer:
            x_norm = layer['norm1'](x)
            # 多头自注意力层，输入需要转换维度为 (seq_len, batch_size, embedding_dim)
            attn_output, _ = layer['attn'](
                x_norm.transpose(0, 1), x_norm.transpose(0, 1), x_norm.transpose(0, 1)
            )
            attn_output = attn_output.transpose(0, 1)  # 转换回形状 (batch_size, seq_len, embedding_dim)
            x = attn_output + x  # 残差连接
            x_norm = layer['norm2'](x)
            x = layer['mlp'](x_norm) + x  # 残差连接

        # 分类头，取 CLS Token 对应的输出进行分类
        x = x[:, 0]
        logits = self.mlp_head(x)
        return logits

# 测试集成的模型
if __name__ == '__main__':
    # MACOS
    # 自动检测设备
    DEVICE = 'mps' if torch.backends.mps.is_available() else 'cpu'  # 支持 MPS 则使用，否则使用 CPU
    print(f"Using device: {DEVICE}")
    # windows
    # DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'  # 设备
    model = SimpleViT().to(DEVICE)  # 将模型移动到指定设备
    input = torch.randn(1, 3, 224, 224).to(DEVICE)  # 生成随机输入数据
    output = model(input)  # 前向传播，获取输出
    print(input.size(), output.size())
    print(f"输出形状: {output.shape}")  # 打印输出的形状，应该为 (1, 1000)

